import jieba  # 分词
from matplotlib import pyplot as plt  # 绘图，数据可视化
from wordcloud import WordCloud  # 词云
from PIL import Image  # 图片处理
import numpy as np  # 矩阵运算
import sqlite3


def fetch_text(db_path):
    """从数据库中获取电影介绍文本"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    sql = 'select introduction from movie'
    data = cursor.execute(sql)
    text = ''.join(item[0] for item in data if item[0])
    cursor.close()
    conn.close()
    return text


def create_word_cloud(text, mask_path, font_path, output_path):
    """生成词云并保存为图片"""
    # 分词
    words = ' '.join(jieba.cut(text))

    # 加载背景图片
    mask = np.array(Image.open(mask_path))

    # 创建词云对象
    wc = WordCloud(
        background_color='white',
        mask=mask,
        font_path=font_path
    )
    wc.generate_from_text(words)

    # 绘制词云图片
    plt.figure()
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')

    # 保存词云图片
    plt.savefig(output_path, dpi=500)
    print(f"词云图片已保存至 {output_path}")


def main():
    db_path = '../影片Top.db'
    mask_path = '../static/img/词云背景图.jpg'
    font_path = 'msyh.ttc'
    output_path = '../static/img/word.jpg'

    # 获取电影介绍文本
    text = fetch_text(db_path)

    # 生成词云
    create_word_cloud(text, mask_path, font_path, output_path)


if __name__ == '__main__':
    main()